import os
import xml.etree.ElementTree as ET
import matplotlib.pyplot as plt
from matplotlib.patches import PathPatch
from matplotlib.path import Path as MplPath
import re

def parse_rgb(style):
    match = re.search(r'fill:rgb\(([\d.%]+),([\d.%]+),([\d.%]+)\)', style)
    if match:
        r, g, b = match.groups()
        return tuple(float(c.strip('%')) / 100 if '%' in c else float(c) / 255 for c in (r, g, b))
    match_hex = re.search(r'fill:\s*#([0-9a-fA-F]{6})', style)
    if match_hex:
        hex_color = match_hex.group(1)
        return tuple(int(hex_color[i:i+2], 16) / 255 for i in (0, 2, 4))
    return (0, 0, 0)  # 默认黑色

def process_svg(svg_file, output_folder):
    with open(svg_file, 'r', encoding='utf-8') as f:
        content = f.read()

    # 解析 SVG 尺寸
    width_match = re.search(r'width="([\d.]+)(?:pt|px)?"', content)
    height_match = re.search(r'height="([\d.]+)(?:pt|px)?"', content)
    if not width_match or not height_match:
        print(f"❌ 无法解析宽高: {svg_file}")
        return

    svg_width = float(width_match.group(1))
    svg_height = float(height_match.group(1))

    fig, ax = plt.subplots(figsize=(svg_width, svg_height))
    ax.set_aspect('equal')
    ax.set_facecolor('white')

    tree = ET.parse(svg_file)
    root = tree.getroot()

    for elem in root.iter('{http://www.w3.org/2000/svg}path'):
        d = elem.attrib.get('d', '')
        style = elem.attrib.get('style', "")
        fill_color = parse_rgb(style)

        # ✅ 只保留包含 'Z' / 'z' 或 'L' 的路径（排除曲线）
        if not any(c in d for c in ['Z', 'z', 'L']):
            continue

        # 提取路径数据
        path_data = re.findall(r'[A-Za-z]|-?\d*\.?\d+', d)
        if not path_data or path_data[0] != 'M':
            continue

        verts = []
        codes = []
        i = 0
        while i < len(path_data):
            cmd = path_data[i]
            i += 1
            if cmd == 'M':
                x = float(path_data[i])
                y = float(path_data[i + 1])
                verts.append((x, y))
                codes.append(MplPath.MOVETO)
                i += 2
            elif cmd == 'L':
                x = float(path_data[i])
                y = float(path_data[i + 1])
                verts.append((x, y))
                codes.append(MplPath.LINETO)
                i += 2
            elif cmd in ('Z', 'z'):
                verts.append((0, 0))  # 结束路径（占位）
                codes.append(MplPath.CLOSEPOLY)

        if verts:
            path = MplPath(verts, codes)
            patch = PathPatch(path, facecolor=fill_color, edgecolor='black')
            ax.add_patch(patch)

    ax.autoscale_view()
    plt.gca().invert_yaxis()
    plt.axis('off')

    output_file = os.path.join(output_folder, os.path.basename(svg_file).replace('.svg', '.png'))
    plt.savefig(output_file, dpi=10, bbox_inches=None, pad_inches=0)
    plt.close()
    print(f"✅ 矩形图像已保存: {output_file}")

def batch_process_svg(input_folder, output_folder):
    if not os.path.exists(output_folder):
        os.makedirs(output_folder)
    for filename in os.listdir(input_folder):
        if filename.lower().endswith('.svg'):
            process_svg(os.path.join(input_folder, filename), output_folder)

# 替换为你的路径
input_folder = r"C:\Users\cenkai\Desktop\PCB_Detection\svg"
output_folder = r"C:\Users\cenkai\Desktop\PCB_Detection\rect_output"

batch_process_svg(input_folder, output_folder)
